{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6f678e62-7bcb-4405-86ae-dce94f494303",
   "metadata": {},
   "source": [
    "# 主要数据加载流程总结"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "070000fc-a7b7-4c56-a2c0-a938d413a790",
   "metadata": {},
   "source": [
    " 完整的章节代码位于[ch02.ipynb](./ch02.ipynb)。\n",
    "\n",
    "这个笔记本包含了主要的要点，即不包含中间步骤的数据加载流程。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "93804da5-372b-45ff-9ef4-8398ba1dd78e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch version: 2.2.1\n",
      "tiktoken version: 0.6.0\n"
     ]
    }
   ],
   "source": [
    "from importlib.metadata import version\n",
    "\n",
    "import tiktoken\n",
    "import torch\n",
    "\n",
    "print(\"torch version:\", version(\"torch\"))\n",
    "print(\"tiktoken version:\", version(\"tiktoken\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0ed4b7db-3b47-4fd3-a4a6-5f4ed5dd166e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tiktoken\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "\n",
    "class GPTDatasetV1(Dataset):\n",
    "    def __init__(self, txt, tokenizer, max_length, stride):\n",
    "        self.tokenizer = tokenizer\n",
    "        self.input_ids = []\n",
    "        self.target_ids = []\n",
    "\n",
    "        # 对全部文本进行分词\n",
    "        token_ids = tokenizer.encode(txt, allowed_special={'<|endoftext|>'})\n",
    "\n",
    "        # 使用滑动窗口将图书分块为最大长度的重叠序列\n",
    "        for i in range(0, len(token_ids) - max_length, stride):\n",
    "            input_chunk = token_ids[i:i + max_length]\n",
    "            target_chunk = token_ids[i + 1: i + max_length + 1]\n",
    "            self.input_ids.append(torch.tensor(input_chunk))\n",
    "            self.target_ids.append(torch.tensor(target_chunk))\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.input_ids)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.input_ids[idx], self.target_ids[idx]\n",
    "\n",
    "\n",
    "def create_dataloader_v1(txt, batch_size=4, max_length=256, \n",
    "                         stride=128, shuffle=True, drop_last=True):\n",
    "    # 分词器初始化\n",
    "    tokenizer = tiktoken.get_encoding(\"gpt2\")\n",
    "\n",
    "    # 创建数据集\n",
    "    dataset = GPTDatasetV1(txt, tokenizer, max_length, stride)\n",
    "\n",
    "    # 创建加载器\n",
    "    dataloader = DataLoader(\n",
    "        dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last)\n",
    "\n",
    "    return dataloader\n",
    "\n",
    "\n",
    "with open(\"the-verdict.txt\", \"r\", encoding=\"utf-8\") as f:\n",
    "    raw_text = f.read()\n",
    "\n",
    "tokenizer = tiktoken.get_encoding(\"gpt2\")\n",
    "encoded_text = tokenizer.encode(raw_text)\n",
    "\n",
    "vocab_size = 50257\n",
    "output_dim = 256\n",
    "block_size = 1024\n",
    "\n",
    "\n",
    "token_embedding_layer = torch.nn.Embedding(vocab_size, output_dim)\n",
    "pos_embedding_layer = torch.nn.Embedding(block_size, output_dim)\n",
    "\n",
    "max_length = 4\n",
    "dataloader = create_dataloader_v1(raw_text, batch_size=8, max_length=max_length, stride=max_length)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "664397bc-6daa-4b88-90aa-e8fc1fbd5846",
   "metadata": {},
   "outputs": [],
   "source": [
    "for batch in dataloader:\n",
    "    x, y = batch\n",
    "\n",
    "    token_embeddings = token_embedding_layer(x)\n",
    "    pos_embeddings = pos_embedding_layer(torch.arange(max_length))\n",
    "\n",
    "    input_embeddings = token_embeddings + pos_embeddings\n",
    "\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d3664332-e6bb-447e-8b96-203aafde8b24",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([8, 4, 256])\n"
     ]
    }
   ],
   "source": [
    "print(input_embeddings.shape)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
